import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from chinese_calendar import is_workday, is_holiday


def schedule_tasks(start_date, extra_workdays, excel_path):
    # 读取Excel文件
    df = pd.read_excel(excel_path)
    df['人/天'] = df['人/天'].apply(lambda x: np.ceil(x) if isinstance(x, float) else x)
    # 将日期字符串转换为datetime对象
    start_date = datetime.strptime(start_date, '%Y-%m-%d')

    # 不管是否为工作日，非工作日的日期会被视为额外工作日
    extra_workdays = set(datetime.strptime(day, '%Y-%m-%d') for day in extra_workdays)

    # 初始化任务开始日期和结束日期列表
    task_start_dates = []
    task_end_dates = []
    not_workdays = []

    # 根据优先级排序，保留原始索引
    sorted_df = df.sort_values(by='优先级').copy()
    original_index = df.index
    sorted_index = sorted_df.index

    current_date = start_date

    for _, row in sorted_df.iterrows():
        # 开始日期不能为非工作日或者额外工作日
        while not is_workday(current_date) and start_date not in extra_workdays:
            current_date += timedelta(days=1)
        people_days = row['人/天']
        task_start_date = current_date
        task_end_date = current_date
        not_workday_str = ''
        people_days_flag = 0

        while True:
            if is_workday(task_end_date) or task_end_date in extra_workdays:
                people_days_flag += 1
                if not is_workday(task_end_date):
                    if not_workday_str:
                        not_workday_str += ','
                    not_workday_str += task_end_date.strftime('%Y-%m-%d')
            if people_days_flag == people_days:
                break
            task_end_date += timedelta(days=1)

        current_date = task_end_date + timedelta(days=1)
        # 添加到列表
        task_start_dates.append(task_start_date.strftime('%Y-%m-%d'))
        task_end_dates.append(task_end_date.strftime('%Y-%m-%d'))
        not_workdays.append(not_workday_str)

    # 使用原始索引更新DataFrame
    df.loc[original_index, '任务开始日期'] = df.loc[sorted_index, '任务开始日期'] = task_start_dates
    df.loc[original_index, '任务结束日期'] = df.loc[sorted_index, '任务结束日期'] = task_end_dates
    df.loc[original_index, '非工作日(可调休)'] = df.loc[sorted_index, '非工作日(可调休)'] = not_workdays

    # 保存到新的Excel文件
    df.to_excel('/Users/eric/Desktop/3.xlsx', index=False)


# 使用示例
start_date = '2024-12-28'
extra_workdays = ['2025-1-18', '2025-1-11']  # 示例额外工作日
excel_path = '/Users/eric/Desktop/1.xlsx'  # Excel文件路径

schedule_tasks(start_date, extra_workdays, excel_path)